# from OCC.Extend.DataExchange import STEPCAFControl_Reader,STEPControl_Reader
# from OCC.Extend.DataExchange import IFSelect_ItemsByEntity,IFSelect_RetDone
# import sys
# from OCC.Core.TDF import TDF_Label
# def read_STEP_file(filename):
#     """ read the STEP file and returns a compound
#     """
#     # 生成一个step模型类
#     reader = STEPControl_Reader()
#     # 加载一个文件并且返回一个状态枚举值
#     status = reader.ReadFile(filename)
#
#     # 如果正常执行且有模型
#     if status == IFSelect_RetDone:  # check status
#         failsonly = False
#         # 如果存在无效或者不完整步骤实体，会显示错误信息
#         reader.PrintCheckLoad(failsonly, IFSelect_ItemsByEntity)
#         reader.PrintCheckTransfer(failsonly, IFSelect_ItemsByEntity)
#
#         # 执行步骤文件转换
#         ok = reader.TransferRoot(1)
#         # 返回生成形状的数量
#         _nbs = reader.NbShapes()
#         # 返回转换后的形状
#         aResShape = reader.Shape(1)
#     else:
#         print("Error: can't read file.")
#         sys.exit(0)
#     return aResShape
#
# read_STEP_file("new_Air.stp")
##Copyright 2010-2014 Thomas Paviot (tpaviot@gmail.com)
##
##This file is part of pythonOCC.
##
##pythonOCC is free software: you can redistribute it and/or modify
##it under the terms of the GNU Lesser General Public License as published by
##the Free Software Foundation, either version 3 of the License, or
##(at your option) any later version.
##
##pythonOCC is distributed in the hope that it will be useful,
##but WITHOUT ANY WARRANTY; without even the implied warranty of
##MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
##GNU Lesser General Public License for more details.
##
##You should have received a copy of the GNU Lesser General Public License
##along with pythonOCC.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function


from OCC.Core.TCollection import TCollection_ExtendedString

from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.XCAFDoc import (XCAFDoc_DocumentTool_ShapeTool,
                              XCAFDoc_DocumentTool_ColorTool,
                              XCAFDoc_DocumentTool_LayerTool,
                              XCAFDoc_DocumentTool_MaterialTool)
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.TDF import TDF_LabelSequence

from OCC.Display.SimpleGui import init_display
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB

from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform

filename = 'new_Air3.stp'
_shapes = []
output_shapes = {}
# create an handle to a document
doc = TDocStd_Document(TCollection_ExtendedString("pythonocc-doc"))

# Get root assembly
shape_tool = XCAFDoc_DocumentTool_ShapeTool(doc.Main())
color_tool  = XCAFDoc_DocumentTool_ColorTool(doc.Main())
l_layers = XCAFDoc_DocumentTool_LayerTool(doc.Main())
l_materials = XCAFDoc_DocumentTool_MaterialTool(doc.Main())

step_reader = STEPCAFControl_Reader()
step_reader.SetColorMode(True)
step_reader.SetLayerMode(True)
step_reader.SetNameMode(True)
step_reader.SetMatMode(True)

status = step_reader.ReadFile(filename)
if status == IFSelect_RetDone:
    step_reader.Transfer(doc)

locs = []


def _get_sub_shapes(lab, loc):
    # global cnt, lvl
    # cnt += 1
    # print("\n[%d] level %d, handling LABEL %s\n" % (cnt, lvl, _get_label_name(lab)))
    # print()
    # print(lab.DumpToString())
    # print()
    # print("Is Assembly    :", shape_tool.IsAssembly(lab))
    # print("Is Free        :", shape_tool.IsFree(lab))
    # print("Is Shape       :", shape_tool.IsShape(lab))
    # print("Is Compound    :", shape_tool.IsCompound(lab))
    # print("Is Component   :", shape_tool.IsComponent(lab))
    # print("Is SimpleShape :", shape_tool.IsSimpleShape(lab))
    # print("Is Reference   :", shape_tool.IsReference(lab))

    # users = TDF_LabelSequence()
    # users_cnt = shape_tool.GetUsers(lab, users)
    # print("Nr Users       :", users_cnt)

    l_subss = TDF_LabelSequence()
    shape_tool.GetSubShapes(lab, l_subss)
    # print("Nb subshapes   :", l_subss.Length())
    l_comps = TDF_LabelSequence()
    shape_tool.GetComponents(lab, l_comps)
    # print("Nb components  :", l_comps.Length())
    # print()
    name = lab.GetLabelName()
    print("Name :", name)

    if shape_tool.IsAssembly(lab):
        l_c = TDF_LabelSequence()
        shape_tool.GetComponents(lab, l_c)
        for i in range(l_c.Length()):
            label = l_c.Value(i + 1)
            if shape_tool.IsReference(label):
                # print("\n########  reference label :", label)
                label_reference = TDF_Label()
                shape_tool.GetReferredShape(label, label_reference)
                loc = shape_tool.GetLocation(label)
                # print("    loc          :", loc)
                # trans = loc.Transformation()
                # print("    tran form    :", trans.Form())
                # rot = trans.GetRotation()
                # print("    rotation     :", rot)
                # print("    X            :", rot.X())
                # print("    Y            :", rot.Y())
                # print("    Z            :", rot.Z())
                # print("    W            :", rot.W())
                # tran = trans.TranslationPart()
                # print("    translation  :", tran)
                # print("    X            :", tran.X())
                # print("    Y            :", tran.Y())
                # print("    Z            :", tran.Z())

                locs.append(loc)
                # print(">>>>")
                # lvl += 1
                _get_sub_shapes(label_reference, loc)
                # lvl -= 1
                # print("<<<<")
                locs.pop()

    elif shape_tool.IsSimpleShape(lab):
        # print("\n########  simpleshape label :", lab)
        shape = shape_tool.GetShape(lab)
        name = lab.GetLabelName()
        print(name)
        # print("    all ass locs   :", locs)

        loc = TopLoc_Location()
        for l in locs:
            # print("    take loc       :", l)
            loc = loc.Multiplied(l)

        # trans = loc.Transformation()
        # print("    FINAL loc    :")
        # print("    tran form    :", trans.Form())
        # rot = trans.GetRotation()
        # print("    rotation     :", rot)
        # print("    X            :", rot.X())
        # print("    Y            :", rot.Y())
        # print("    Z            :", rot.Z())
        # print("    W            :", rot.W())
        # tran = trans.TranslationPart()
        # print("    translation  :", tran)
        # print("    X            :", tran.X())
        # print("    Y            :", tran.Y())
        # print("    Z            :", tran.Z())
        c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
        colorSet = False
        if (color_tool.GetInstanceColor(shape, 0, c) or
                color_tool.GetInstanceColor(shape, 1, c) or
                color_tool.GetInstanceColor(shape, 2, c)):
            color_tool.SetInstanceColor(shape, 0, c)
            color_tool.SetInstanceColor(shape, 1, c)
            color_tool.SetInstanceColor(shape, 2, c)
            colorSet = True
            n = c.Name(c.Red(), c.Green(), c.Blue())
            print('    instance color Name & RGB: ', c, n, c.Red(), c.Green(), c.Blue())

        if not colorSet:
            if (color_tool.GetColor(lab, 0, c) or
                    color_tool.GetColor(lab, 1, c) or
                    color_tool.GetColor(lab, 2, c)):
                color_tool.SetInstanceColor(shape, 0, c)
                color_tool.SetInstanceColor(shape, 1, c)
                color_tool.SetInstanceColor(shape, 2, c)

                n = c.Name(c.Red(), c.Green(), c.Blue())
                print('    shape color Name & RGB: ', c, n, c.Red(), c.Green(), c.Blue())

        shape_disp = BRepBuilderAPI_Transform(shape, loc.Transformation()).Shape()
        if not shape_disp in output_shapes:
            output_shapes[shape_disp] = [lab.GetLabelName(), c]
        for i in range(l_subss.Length()):
            lab_subs = l_subss.Value(i + 1)
            # print("\n########  simpleshape subshape label :", lab)
            shape_sub = shape_tool.GetShape(lab_subs)

            c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
            colorSet = False
            if (color_tool.GetInstanceColor(shape_sub, 0, c) or
                    color_tool.GetInstanceColor(shape_sub, 1, c) or
                    color_tool.GetInstanceColor(shape_sub, 2, c)):
                color_tool.SetInstanceColor(shape_sub, 0, c)
                color_tool.SetInstanceColor(shape_sub, 1, c)
                color_tool.SetInstanceColor(shape_sub, 2, c)
                colorSet = True
                n = c.Name(c.Red(), c.Green(), c.Blue())
                print('    instance color Name & RGB: ', c, n, c.Red(), c.Green(), c.Blue())

            if not colorSet:
                if (color_tool.GetColor(lab_subs, 0, c) or
                        color_tool.GetColor(lab_subs, 1, c) or
                        color_tool.GetColor(lab_subs, 2, c)):
                    color_tool.SetInstanceColor(shape, 0, c)
                    color_tool.SetInstanceColor(shape, 1, c)
                    color_tool.SetInstanceColor(shape, 2, c)

                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    print('    shape color Name & RGB: ', c, n, c.Red(), c.Green(), c.Blue())
            shape_to_disp = BRepBuilderAPI_Transform(shape_sub, loc.Transformation()).Shape()
            # position the subshape to display
            if not shape_to_disp in output_shapes:
                output_shapes[shape_to_disp] = [lab_subs.GetLabelName(), c]

    elif shape_tool.IsShape(lab):
        name = lab.GetLabelName()
        print(name)




def _get_shapes():
    labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(labels)
    # global cnt
    # cnt += 1

    print()
    print("Number of shapes at root :", labels.Length())
    print()
    for i in range(labels.Length()):
        root_item = labels.Value(i + 1)
        _get_sub_shapes(root_item, None)


_get_shapes()

#
# Display
#
display, start_display, add_menu, add_function_to_menu = init_display("wx")
display.DisplayShape(_shapes, update=True)
start_display()
